{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "522c2f1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ef803128",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using matplotlib backend: Qt5Agg\n"
     ]
    }
   ],
   "source": [
    "%matplotlib\n",
    "xBegin = 0\n",
    "xend = 1\n",
    "D = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fbec1df6",
   "metadata": {},
   "outputs": [],
   "source": [
    "bottom = 0\n",
    "top = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d6d44a42",
   "metadata": {},
   "outputs": [],
   "source": [
    "def uBegin(x,x0 = 0.5):\n",
    "    return - 4 * (x - x0) ** 2 + 1 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3142e16d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x23fde5bcf70>]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = np.linspace(0,1,100)\n",
    "u0 = uBegin(x)\n",
    "plt.plot(x,u0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "91e375c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "dt = 0.002\n",
    "dx = 0.05\n",
    "signa = D * dt / (dx ** 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "378c05b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.7999999999999998"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "signa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "d63f9d6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "xArr = np.arange(xBegin,xend + dx,dx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "ff121a60",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 100\n",
    "tBegin = 0\n",
    "tEnd = dt * n\n",
    "tArr = np.arange(tBegin,tEnd,dt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "9c32195c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "21"
      ]
     },
     "execution_count": 79,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xArr.size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "1e19a381",
   "metadata": {},
   "outputs": [],
   "source": [
    "A = np.identity(xArr.size - 2)\n",
    "\n",
    "B = np.roll(A,1)\n",
    "B[0,0] = 0\n",
    "\n",
    "C = np.roll(A,-1)\n",
    "C[-1,-1] = 0\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "id": "9405bee0",
   "metadata": {},
   "outputs": [],
   "source": [
    "AFinal = (1+2*signa)*A - signa * B - signa * C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "id": "22d40b21",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 2.6, -0.8,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,\n",
       "         0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ],\n",
       "       [-0.8,  2.6, -0.8,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,\n",
       "         0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ],\n",
       "       [ 0. , -0.8,  2.6, -0.8,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,\n",
       "         0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ],\n",
       "       [ 0. ,  0. , -0.8,  2.6, -0.8,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,\n",
       "         0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ],\n",
       "       [ 0. ,  0. ,  0. , -0.8,  2.6, -0.8,  0. ,  0. ,  0. ,  0. ,  0. ,\n",
       "         0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ],\n",
       "       [ 0. ,  0. ,  0. ,  0. , -0.8,  2.6, -0.8,  0. ,  0. ,  0. ,  0. ,\n",
       "         0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ],\n",
       "       [ 0. ,  0. ,  0. ,  0. ,  0. , -0.8,  2.6, -0.8,  0. ,  0. ,  0. ,\n",
       "         0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ],\n",
       "       [ 0. ,  0. ,  0. ,  0. ,  0. ,  0. , -0.8,  2.6, -0.8,  0. ,  0. ,\n",
       "         0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ],\n",
       "       [ 0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. , -0.8,  2.6, -0.8,  0. ,\n",
       "         0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ],\n",
       "       [ 0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. , -0.8,  2.6, -0.8,\n",
       "         0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ],\n",
       "       [ 0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. , -0.8,  2.6,\n",
       "        -0.8,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ],\n",
       "       [ 0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. , -0.8,\n",
       "         2.6, -0.8,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ],\n",
       "       [ 0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,\n",
       "        -0.8,  2.6, -0.8,  0. ,  0. ,  0. ,  0. ,  0. ],\n",
       "       [ 0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,\n",
       "         0. , -0.8,  2.6, -0.8,  0. ,  0. ,  0. ,  0. ],\n",
       "       [ 0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,\n",
       "         0. ,  0. , -0.8,  2.6, -0.8,  0. ,  0. ,  0. ],\n",
       "       [ 0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,\n",
       "         0. ,  0. ,  0. , -0.8,  2.6, -0.8,  0. ,  0. ],\n",
       "       [ 0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,\n",
       "         0. ,  0. ,  0. ,  0. , -0.8,  2.6, -0.8,  0. ],\n",
       "       [ 0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,\n",
       "         0. ,  0. ,  0. ,  0. ,  0. , -0.8,  2.6, -0.8],\n",
       "       [ 0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,  0. ,\n",
       "         0. ,  0. ,  0. ,  0. ,  0. ,  0. , -0.8,  2.6]])"
      ]
     },
     "execution_count": 131,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "AFinal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 132,
   "id": "e164ea4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "bVector = np.zeros(xArr.size - 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 133,
   "id": "580c9f8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "bVector[0] = signa * bottom\n",
    "bVector[-1] = signa * top"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "id": "5294b868",
   "metadata": {},
   "outputs": [],
   "source": [
    "AInv = np.linalg.inv(AFinal)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 135,
   "id": "2b0c21bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "uCurrent = uBegin(xArr[1:-1])\n",
    "plt.plot(xArr[1:-1],uCurrent)\n",
    "uArr = np.zeros((xArr.size - 2,tArr.size))\n",
    "uArr[:,0] = uCurrent\n",
    "for i in range(1,n):\n",
    "    uNext = AInv @ uCurrent + AInv @ bVector\n",
    "    uArr[:,i] = uNext\n",
    "    plt.plot(xArr[1:-1],uNext)\n",
    "    uCurrent = uNext\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "id": "ef270bb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "XX, YY = np.meshgrid(xArr[1:-1],tArr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 137,
   "id": "a1c94f6d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(100, 19)"
      ]
     },
     "execution_count": 137,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "XX.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "id": "ecd5704f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(100, 19)"
      ]
     },
     "execution_count": 138,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "YY.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "id": "63ad9ff3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using matplotlib backend: Qt5Agg\n"
     ]
    }
   ],
   "source": [
    "%matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "id": "a392e430",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<mpl_toolkits.mplot3d.art3d.Poly3DCollection at 0x23fe4a7c6d0>"
      ]
     },
     "execution_count": 144,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fig,ax = plt.subplots(subplot_kw={\"projection\":\"3d\"})\n",
    "ax.plot_surface(XX,YY,uArr.T,cmap = plt.cm.spring)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db381cb1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25a25cdb",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
